import datetime

# 硬编码2025年国定假日
NATIONAL_HOLIDAYS = {
    # 元旦（1天）
    datetime.date(2025, 1, 1),
    
    # 春节（8天）
    datetime.date(2025, 1, 28), datetime.date(2025, 1, 29),
    datetime.date(2025, 1, 30), datetime.date(2025, 1, 31),
    datetime.date(2025, 2, 1), datetime.date(2025, 2, 2),
    datetime.date(2025, 2, 3), datetime.date(2025, 2, 4),
    
    # 清明节（3天）
    datetime.date(2025, 4, 4), datetime.date(2025, 4, 5),
    datetime.date(2025, 4, 6),
    
    # 劳动节（5天）
    datetime.date(2025, 5, 1), datetime.date(2025, 5, 2),
    datetime.date(2025, 5, 3), datetime.date(2025, 5, 4),
    datetime.date(2025, 5, 5),
    
    # 端午节（3天）
    datetime.date(2025, 5, 31), datetime.date(2025, 6, 1),
    datetime.date(2025, 6, 2),
    
    # 中秋+国庆（8天）
    datetime.date(2025, 10, 1), datetime.date(2025, 10, 2),
    datetime.date(2025, 10, 3), datetime.date(2025, 10, 4),
    datetime.date(2025, 10, 5), datetime.date(2025, 10, 6),
    datetime.date(2025, 10, 7), datetime.date(2025, 10, 8)
}

# ================== 2025年调休工作日硬编码数据 ==================
WORKDAYS = {
    datetime.date(2025, 1, 26),  # 周日补班（春节前）
    datetime.date(2025, 2, 8),   # 周六补班（春节后）
    datetime.date(2025, 4, 27),  # 周日补班（劳动节前）
    datetime.date(2025, 9, 28),  # 周日补班（国庆前）
    datetime.date(2025, 10, 11)  # 周六补班（国庆后）
}


def is_workday(date):
    """
    判断给定日期是否为工作日
    优化策略：使用集合查询O(1)时间复杂度 + 优化判断顺序
    """
    if date in NATIONAL_HOLIDAYS:
        return False
    if date in WORKDAYS:
        return True
    return date.weekday() < 5  # 直接返回是否平日

def next_n_workdays(start_date: datetime.date, n: int) -> list:
    """
    计算从start_date的下一个工作日开始的n个工作日的日期
    :param start_date: 起始日期（datetime.date对象）
    :param n: 工作日数量（正整数）
    :return: 包含n个工作日的日期列表
    """
    workdays_list = []
    current_date = start_date + datetime.timedelta(days=1)  # 从下一天开始
    while len(workdays_list) < n:
        if is_workday(current_date):
            workdays_list.append(current_date)
        current_date += datetime.timedelta(days=1)
    return workdays_list

def count_workdays(start_date, end_date):
    """
    计算给定的两个日期之间共有几个工作日
    优化点：使用生成器表达式 + 提前终止循环（当current_date超过end_date时）
    """
    return sum(1 for current_date in date_generator(start_date, end_date) if is_workday(current_date))

def date_generator(start, end):
    """日期生成器，减少内存占用"""
    current = start
    while current <= end:
        yield current
        current += datetime.timedelta(days=1)

# 示例保持相同
# print(is_workday(datetime.date(2024, 1, 1)))  # False
# print(next_n_workdays(datetime.date(2024, 1, 1), 5))
# print(count_workdays(datetime.date(2024, 1, 1), datetime.date(2024, 1, 7)))  # 4


def calculate_deadline(start_time: datetime.datetime, n: int) -> datetime.datetime:
    """
    计算n个工作日后的截止时间（每日8:30-17:00为工作时间）
    
    参数：
    start_time -- 起始时间（datetime对象）
    n -- 工作日天数（正整数）
    
    返回：
    datetime.datetime -- 截止时间
    """
    if n <= 0:
        raise ValueError("n必须为正整数")

    # 调整开始时间为有效工作时间
    adjusted_start = adjustLastWorkTime(start_time)

    # 计算n个工作日后的日期（从adjusted_start的下一天开始）
    target_date = next_n_workdays(adjusted_start.date(), n)[-1]

    # 拼接时间并返回
    deadline = datetime.datetime.combine(target_date, adjusted_start.time())
    return adjustLastWorkTime(deadline)


def adjustLastWorkTime(dt: datetime.datetime) -> datetime.datetime:
    """查找最近工作日的17点（含当天可能）"""
    current_date = dt.date()
    # 先检查当天是否是工作日且时间已过17点
    if is_workday(current_date):
        if dt.time() >= datetime.time(8, 30) and dt.time() <= datetime.time(17, 0):
            return dt
        elif dt.time() > datetime.time(17, 0):
            return dt.replace(hour=17, minute=0, second=0, microsecond=0)

    # 向前查找最近的工作日
    while True:
        current_date -= datetime.timedelta(days=1)
        if is_workday(current_date):
            return datetime.datetime.combine(current_date, datetime.time(17, 0))


def test():
    print(calculate_deadline(datetime.datetime(2025, 1, 2, 19, 0), 2))
    print(calculate_deadline(datetime.datetime(2025, 1, 2, 13, 0), 2))
    print(calculate_deadline(datetime.datetime(2025, 2, 4, 18, 0), 2))
    print(calculate_deadline(datetime.datetime(2025, 10, 14, 20, 0), 3))
    print(calculate_deadline(datetime.datetime(2025, 10, 14, 12, 0), 7))


if __name__ == "__main__":
    test()